import math


class ClipGrad:
    def __init__(self, max_norm):
        """
        梯度裁剪，用于防止梯度爆炸。

        Args:
            max_norm (float): 允许的梯度最大范数。
        """
        self.max_norm = max_norm

    def __call__(self, params):
        """
        对每个参数的梯度进行裁剪，使其范数不超过指定的最大值。

        Args:
            params (list): 参数列表。

        Returns:
            None
        """
        total_norm = 0
        for param in params:
            total_norm += (param.grad.data ** 2).sum()
        total_norm = math.sqrt(float(total_norm))

        rate = self.max_norm / (total_norm + 1e-6)
        if rate < 1:
            for param in params:
                param.grad.data *= rate


